# -*- coding: utf-8 -*-
"""
Created on Sun Oct 11 10:09:35 2020
用于双量子点量子模型
单量子比特的归零操作
多个初始态对多个目标态的归零保真度

@author: Waikikilick
"""

import numpy as np
from scipy.linalg import expm
from time import *
import multiprocessing as mp
np.random.seed(1)
T = np.pi
dt = np.pi/20
step_max = T/dt
sx = np.mat([[0, 1], [1, 0]], dtype=complex)
sy = np.mat([[0, -1j], [1j, 0]], dtype=complex)
sz = np.mat([[1, 0], [0, -1]], dtype=complex)

# a0,a1,a2,a3,a4,a5,a6,a7,a8,a9 = 0,0,0,0,0,0,0,0,0,0 #统计各动作被选用的频率

action_space = np.mat([[1,0,0], #可以选择的动作范围，各列的每项分别代表着 sigma x, y, z 前面的系数。
                       [2,0,0], #每次执行的动作都是单独的绕 x, y, z 轴一定角度的旋转
                       [0,1,0], # x, y 方向的值可以取负，但 z 方向的只能取正值
                       [0,2,0],
                       [0,0,1],
                       [0,0,2],
                       [-1,0,0],
                       [-2,0,0],
                       [0,-1,0],
                       [0,-2,0]])

theta_num = 6 #除了 0 和 Pi 两个点之外，点的数量
varphi_num = 21#varphi 角度一圈上的点数

theta = np.linspace(0,np.pi,theta_num+2,endpoint=True) 
varphi = np.linspace(0,np.pi*2,varphi_num,endpoint=False) 

def psi_set():
    psi_set = []
    for ii in range(1,theta_num+1):
        for jj in range(varphi_num):
            psi_set.append(np.mat([[np.cos(theta[ii]/2)],[np.sin(theta[ii]/2)*(np.cos(varphi[jj])+np.sin(varphi[jj])*(0+1j))]]))
    psi_set.append(np.mat([[1], [0]], dtype=complex))
    psi_set.append(np.mat([[0], [1]], dtype=complex))
    return psi_set

target_set = psi_set()
init_set = psi_set()

#动作直接选最优的
def step(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        H = float(action_space[action,0])*sx/2 + float(action_space[action,1])*sy/2 - float(action_space[action,2])*sz/2
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        psi_list.append(psi_)
        fid_list.append(fid)
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    psi_ = psi_list[best_action]
    return best_action, best_fid, psi_

#动作选最优的，或者最差的
def step1(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        H = float(action_space[action,0])*sx/2 + float(action_space[action,1])*sy/2 - float(action_space[action,2])*sz/2
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
    
    if F < max(fid_list):
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    else:
        best_action = fid_list.index(min(fid_list))
        best_fid = min(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_

#动作选最优的，或者次优的
def step2(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        H = float(action_space[action,0])*sx/2 + float(action_space[action,1])*sy/2 - float(action_space[action,2])*sz/2
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
        
    if F < max(fid_list):
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    else:
        del action_list[fid_list.index(max(fid_list))]
        del psi_list[fid_list.index(max(fid_list))]
        del fid_list[fid_list.index(max(fid_list))]
        
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
        
    psi_ = psi_list[best_action]
    return best_action, best_fid, psi_

def job(target_psi):
    fid_list, time_list = [], []
    for psi1 in init_set:
        start_time = time()
        psi = psi1
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        fid_max = 0
        fid_max1 = 0
        fid_max2 = 0
        fid_max0 = 0
        
        step_n = 0
        start_time = time()
        while True:
            action, F, psi_ = step1(psi,target_psi,F)
            fid_max1 = max(F,fid_max1)
            psi = psi_
            step_n += 1
            if fid_max1>0.999 or step_n>step_max:
                break
        end_time = time()
        time_list.append(end_time-start_time) 
        
        step_n = 0
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        psi = psi1
        start_time = time()
        while True:
            action, F, psi_ = step2(psi,target_psi,F)
            fid_max2 = max(F,fid_max2)
            psi = psi_
            step_n += 1
            if fid_max2>0.999 or step_n>step_max:
                break 
        end_time = time()
        time_list.append(end_time-start_time)   
        fid_max = max(fid_max1,fid_max2)
        
    
        step_n = 0
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        psi = psi1
        start_time = time()
        while True:
            action, F, psi_ = step(psi,target_psi,F)
            fid_max0 = max(F,fid_max0)
            psi = psi_
            step_n += 1
            if fid_max0>0.999 or step_n>step_max:
                break 
        fid_max = max(fid_max,fid_max0)  
        end_time = time()
        time_list.append(end_time-start_time)
        fid_list.append(fid_max)
    return  (np.mean(time_list),np.mean(fid_list))

def multicore():
    pool = mp.Pool()
    data = pool.map(job, target_set)
    return data
    

if __name__ == '__main__':
    
    time1 = time()
    data = multicore()
    data.sort()
    # print(data)
    time2 = time()
    # print(np.mean(F_list))
    print('time_cost is: ',time2-time1)
    print(data)

# [(0.03544040955603123, 0.9995540107497188), (0.036607105284929276, 0.9995540107497185), (0.03737686139841875, 0.9995125756924361), (0.037506623193621635, 0.9995366303277342), (0.03870131137470404, 0.9995340631758385), (0.03922828038533529, 0.9994607633414898), (0.03927763178944588, 0.9995460533336459), (0.039528580382466316, 0.9995188853013959), (0.03967664328714212, 0.999507929184891), (0.03988125051061312, 0.9995366303277345), (0.03988455484310786, 0.9994850944224671), (0.039936394120256104, 0.9994904903874462), (0.04009026847779751, 0.9995340631758378), (0.0401108389099439, 0.9994262042817064), (0.04013702956338724, 0.9995306864977332), (0.04048481024801731, 0.9995460533336453), (0.04070428013801575, 0.9995770571775677), (0.04071948118507862, 0.9995065968535366), (0.04081843917568525, 0.9995770571775673), (0.040891220793128014, 0.9995079291848907), (0.040977381790677704, 0.9995381557291694), (0.04099689982831478, 0.9995381557291688), (0.041225723922252655, 0.9995306864977327), (0.041262704879045486, 0.9994484890569242), (0.04148705117404461, 0.9993963640685819), (0.04154879413545132, 0.9995188853013954), (0.04159103209773699, 0.9993321384538586), (0.04164401193459829, 0.9994711079783871), (0.041703797255953155, 0.9994774534601343), (0.04178505142529806, 0.9994774534601336), (0.0418143297235171, 0.999389117400597), (0.041817442824443184, 0.9994341030357914), (0.041846247389912605, 0.9994277361138162), (0.04193205013871193, 0.9994850944224668), (0.042105297868450485, 0.9994868779795234), (0.04213308977584044, 0.9995125756924356), (0.042199498042464256, 0.9994847719313824), (0.042302511632442474, 0.999448770389581), (0.042375133062402405, 0.9994201956923214), (0.04237553353110949, 0.9994484890569237), (0.04243923289080461, 0.9994400424435185), (0.04254198260605335, 0.9994262042817057), (0.04255476531883081, 0.999531320438355), (0.04263504780828953, 0.999506596853536), (0.04269216333826383, 0.9994418324776196), (0.04271160997450352, 0.9995312501317186), (0.042748210951685905, 0.9994548054567872), (0.04290020217498144, 0.9994344492569697), (0.04290997236967087, 0.9994607633414894), (0.04297946331401666, 0.9995358246179085), (0.043084532022476196, 0.9995312501317182), (0.043093898023168244, 0.999486877979523), (0.04314608499407768, 0.9994344492569693), (0.04334863523642222, 0.9995147454366144), (0.04350543456772963, 0.9994400424435181), (0.04350596542159716, 0.9994847719313821), (0.043513245259722076, 0.9994297317319815), (0.04354531938831011, 0.9993789657781865), (0.04359667127331098, 0.9994418498151315), (0.043609737108151116, 0.9994548054567868), (0.04371277056634426, 0.9994711079783866), (0.043739273523290954, 0.999573102765684), (0.04376576468348503, 0.9994760814110155), (0.04391111681858698, 0.9994904903874458), (0.04392391194899877, 0.9995206790708228), (0.04425921477377415, 0.9994442339539468), (0.044329605996608734, 0.9994487703895805), (0.04441262471179167, 0.9995147454366143), (0.04444404629369577, 0.999514484401935), (0.04456849520405134, 0.9994094370995575), (0.044649612779418625, 0.9994277361138157), (0.044682445004582405, 0.9995110420974347), (0.044691442201534905, 0.9994418324776193), (0.044748094553748764, 0.9994114725257053), (0.04484459261099497, 0.9995731027656835), (0.0449123860647281, 0.9994094370995572), (0.04497621518870195, 0.9994075732558492), (0.04536209503809611, 0.9994760814110153), (0.04541334634025892, 0.9994823498684329), (0.045653076842427254, 0.9995110420974345), (0.045657375206549965, 0.9994997090832406), (0.0457208106915156, 0.9994341030357909), (0.04575435320536295, 0.9995313204383549), (0.045764063174525894, 0.9994362014931584), (0.04582273215055466, 0.9993891174005963), (0.04582499464352926, 0.999375857177609), (0.04591769973436991, 0.9995358246179081), (0.046415701508522034, 0.9990514218806049), (0.0464768260717392, 0.9994418498151311), (0.046547003711263336, 0.9994075732558487), (0.04663265993197759, 0.9994362014931581), (0.04668626996378104, 0.9995206790708222), (0.04673146021862825, 0.9994823498684327), (0.046797540659705796, 0.9994997090832402), (0.04695195828874906, 0.9993789657781861), (0.04699634946882725, 0.9992626016568968), (0.047139509891470276, 0.9994201956923208), (0.04744847367207209, 0.999318335636674), (0.04747025420268377, 0.9995207069583345), (0.04761122105022272, 0.9994442339539463), (0.04818039139111837, 0.9990079323385297), (0.04875412334998449, 0.9994551053468369), (0.04880924088259538, 0.9995401661775865), (0.048833453406890236, 0.9994551053468372), (0.04895277818044027, 0.9993321384538579), (0.04897471579412619, 0.9993013807849256), (0.04933937142292658, 0.9992999654172497), (0.04979523333410422, 0.9995144844019346), (0.049874912947416306, 0.9993183356366736), (0.050028530259927116, 0.9995777587841987), (0.05089659926791986, 0.9990514218806044), (0.05099104344844818, 0.9994114725257046), (0.051117461174726486, 0.9993963640685813), (0.051301086942354836, 0.9994297317319811), (0.05136891764899095, 0.999520706958334), (0.05170766388376554, 0.9989346144275266), (0.05219467729330063, 0.999007932338529), (0.05269036628305912, 0.999540166177586), (0.05289027219017347, 0.9989399170707955), (0.052925666173299156, 0.9994664195841062), (0.053269789243737854, 0.9995489943208509), (0.053601166854302086, 0.9992975575992266), (0.05381078210969766, 0.9992989411366326), (0.05386141190926234, 0.9992626016568962), (0.054518397276600204, 0.9993758571776086), (0.05496348564823469, 0.9995777587841981), (0.05621222034096718, 0.9994664195841054), (0.05881931632757187, 0.9995489943208504)]


